import numpy as np
from astropy.io import fits
from astropy.table import Table
from astropy.stats import sigma_clipped_stats
import matplotlib.pyplot as plt
from photutils.psf import BasicPSFPhotometry, DAOGroup, IntegratedGaussianPRF
from photutils.background import MMMBackground

class PSFFitter:
    def __init__(self, sigma=2.5, fitshape=(31, 31)):
        self.sigma = sigma
        self.fitshape = fitshape
        self.result = None
        self.data = None
        self.closest_source = None

    def load_fits(self, fits_path):
        self.data = fits.getdata(fits_path)
        return self.data

    def guess_init_position(self):
        """使用图像中心和最大值作为初始猜测"""
        y0, x0 = np.array(self.data.shape) // 2
        init_guesses = Table()
        init_guesses['x_0'] = [x0]
        init_guesses['y_0'] = [y0]
        init_guesses['flux_0'] = [np.max(self.data)]
        return init_guesses

    def run_psf_fit(self, init_guesses=None):
        """执行 PSF 拟合"""
        if init_guesses is None:
            init_guesses = self.guess_init_position()

        psf_model = IntegratedGaussianPRF(sigma=self.sigma)
        psf_model.flux.fixed = False

        photometry = BasicPSFPhotometry(
            group_maker=DAOGroup(2.0),
            bkg_estimator=MMMBackground(),
            psf_model=psf_model,
            fitshape=self.fitshape,
        )

        self.result = photometry.do_photometry(self.data, init_guesses=init_guesses)
        return self.result

    def find_closest_to_center(self):
        """找到拟合结果中最接近图像中心的源"""
        if self.result is None:
            raise RuntimeError("请先运行 run_psf_fit()")

        yc, xc = np.array(self.data.shape) / 2
        distances = np.sqrt((self.result['x_fit'] - xc)**2 + (self.result['y_fit'] - yc)**2)
        closest_idx = np.argmin(distances)

        self.closest_source = {
            'x': float(self.result['x_fit'][closest_idx]),
            'y': float(self.result['y_fit'][closest_idx]),
            'flux': float(self.result['flux_fit'][closest_idx]),
            'distance_to_center': float(distances[closest_idx])
        }
        return self.closest_source

    def plot_result(self, savepath=None):
        """可视化拟合中心"""
        if self.result is None:
            raise RuntimeError("请先运行 run_psf_fit()")

        median = np.median(self.data)
        std = np.std(self.data)

        plt.figure(figsize=(6, 6))
        plt.imshow(self.data, origin='lower', cmap='gray', vmin=median-2*std, vmax=median+8*std)
        plt.plot(self.result['x_fit'], self.result['y_fit'], 'rx', label='PSF fits')

        if self.closest_source is not None:
            plt.plot(self.closest_source['x'], self.closest_source['y'], 'go', label='Closest to center', markersize=10)

        plt.title('PSF Fit Result')
        plt.legend()
        if savepath:
            plt.savefig(savepath, dpi=150)
        else:
            plt.show()

psf_tool = PSFFitter()
psf_tool.load_fits('MTC_J062018.43P711440.50_cutout.fits')
psf_tool.run_psf_fit()
info = psf_tool.find_closest_to_center()
print(info)
psf_tool.plot_result("psf_result.png")
